Переход от режима Элегантного выполнения PyTorch к Triton требует смены подхода: вместо восприятия тензоров как монолитных объектов их следует рассматривать как совокупность дискретных, управляемых блоков или плиток.

1. Тензоры PyTorch и Triton

Очень важно различать тензоры Triton от тензоры PyTorch. Тензор PyTorch — это объект на стороне хоста на языке Python обёртка формы, типа данных, устройства, шагов и метаданных хранилища. В отличие от этого, Triton работает с сырыми указателями на данные внутри конкретных блоков памяти, что позволяет проводить оптимизацию на более низком уровне.

2. Бутылочное горло режима Элегантного выполнения

При стандартном режиме элегантного выполнения каждая операция (например, сложение, затем функция активации) требует отдельной загрузки ядра и обращения к глобальной памяти. Это основное бутылочное горло при современных вычислениях на видеокартах. Triton преодолевает это, путём объединения операций внутри одного ядра, которое обрабатывает блоки данных (например, 128, 256 или 512 элементов) непосредственно в памяти чипа.

3. Парадигма на основе блоков

Вместо скалярного мышления, характерного для потоков CUDA, Triton использует SPMD (единая программа, множество данных) на уровне блоков. Вы пишете одно ядро, и Triton запускает несколько экземпляров по всей сетке. Каждый экземпляр использует свой program_id для расчёта того, какой «кусок» памяти он владеет.

4. Настройка среды

Чтобы начать, установите Triton в чистой среде (с помощью Conda или venv), чтобы избежать конфликта зависимостей с существующими наборами инструментов CUDA: pip install triton.